package kvpair

import (
	"bytes"

	multiproof "github.com/crate-crypto/go-ipa"
	"github.com/crate-crypto/go-ipa/bandersnatch"
	"github.com/crate-crypto/go-ipa/bandersnatch/fr"
	"github.com/crate-crypto/go-ipa/common"
	"github.com/crate-crypto/go-ipa/ipa"
	gthCrypto "github.com/ethereum/go-ethereum/crypto"
)

type KVpair struct {
	Key   byte
	Value [32]byte
}

type Fr = fr.Element
type Point = bandersnatch.PointAffine

type dataRep struct {
	poly   []fr.Element // length shold be 256
	commit Point        // srs poly commient
	conf   *ipa.IPAConfig
}

func NewRep(data []Fr) *dataRep {
	conf := GetConf()
	return &dataRep{
		poly:   data,
		commit: conf.Commit(data),
		conf:   conf,
	}
}

func (s *dataRep) CreateProve(ins []KVpair) *multiproof.MultiProof {
	prover_transcript := common.NewTranscript("multiproof")
	cs := make([]*Point, len(ins))
	fs := make([][]Fr, len(ins))
	zs := make([]uint8, len(ins))
	for i, val := range ins {
		cs[i] = &s.commit
		fs[i] = s.poly
		zs[i] = uint8(val.Key)
	}
	return multiproof.CreateMultiProof(prover_transcript, s.conf, cs, fs, zs)
}

func (s *dataRep) VerifyProve(proof *multiproof.MultiProof, ins []KVpair) bool {
	cs := make([]*Point, len(ins))
	fs := make([][]Fr, len(ins))
	zs := make([]uint8, len(ins))
	ys := make([]*fr.Element, len(ins))
	for i, val := range ins {
		cs[i] = &s.commit
		fs[i] = s.poly
		zs[i] = uint8(val.Key)
		var yfr Fr
		yfr.SetBytes(val.Value[:])
		ys[i] = &yfr
	}
	verifier_transcript := common.NewTranscript("multiproof")
	return multiproof.CheckMultiProof(verifier_transcript, s.conf, proof, cs, ys, zs)
}

// GenPolyAndKVpair generate poly and data kvpair set
func GenPolyAndKVpair(ins [][]byte) ([]Fr, []KVpair) {
	n := len(ins)
	if len(ins) > 256 {
		panic("polynomial cannot exceed 256 coefficients")
	}

	root := make([]fr.Element, 256, 256)
	kvs := make([]KVpair, n)
	for i, one := range ins {
		hash := gthCrypto.Keccak256(one)
		root[i].SetBytes(hash)
		tmp := KVpair{
			Key: byte(i),
		}
		copy(tmp.Value[:], hash)
		kvs[i] = tmp
	}

	for i := n; i < 256; i++ {
		root[i] = fr.Zero()
	}
	return root, kvs
}

// SerializeProof 序列化proof
func SerializeProof(proof multiproof.MultiProof) []byte {
	var buf = new(bytes.Buffer)
	proof.Write(buf)
	return buf.Bytes()
}

func DeserializeProof(data []byte) multiproof.MultiProof {
	var buf = new(bytes.Buffer)
	var got_proof multiproof.MultiProof
	buf.Write(data)
	got_proof.Read(buf)
	return got_proof
}
